package com.tongcao.cn.aspect;

import com.fasterxml.jackson.databind.ObjectMapper;
import com.google.common.util.concurrent.RateLimiter;
import com.tongcao.cn.annotation.RateLimit;
import com.tongcao.cn.dfire.result.ResultMap;
import com.tongcao.cn.dfire.utils.ApiResultUtil;
import com.tongcao.cn.util.FastJsonUtils;
import org.aspectj.lang.ProceedingJoinPoint;
import org.aspectj.lang.Signature;
import org.aspectj.lang.annotation.Around;
import org.aspectj.lang.annotation.Aspect;
import org.aspectj.lang.annotation.Pointcut;
import org.aspectj.lang.reflect.MethodSignature;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.context.annotation.Scope;
import org.springframework.stereotype.Component;

import javax.annotation.Resource;
import javax.servlet.ServletOutputStream;
import javax.servlet.http.HttpServletResponse;
import java.io.IOException;
import java.lang.reflect.Method;
import java.util.Map;
import java.util.concurrent.ConcurrentHashMap;

/**
 * @author:tongcao
 * @since:2019/7/9 19:48
 * Description:
 */

@Component
@Aspect
@Scope
public class RateLimitAspect {


    private Logger log = LoggerFactory.getLogger(this.getClass());


    //用来存放不同接口的RateLimiter(key为接口名称，value为RateLimiter)
    private Map<String, RateLimiter> map = new ConcurrentHashMap<>();


    private static ObjectMapper objectMapper = new ObjectMapper();

    private RateLimiter rateLimiter;

    @Resource
    private HttpServletResponse response;

    @Pointcut("@annotation(com.tongcao.cn.annotation.RateLimit)")
    public void serviceLimit() {
    }

    @Around("serviceLimit()")
    public Object around(ProceedingJoinPoint joinPoint) throws NoSuchMethodException {
        {

            Object obj = null;
            //获取拦截的方法名
            Signature sig = joinPoint.getSignature();
            //获取拦截的方法名
            MethodSignature msig = (MethodSignature) sig;
            //返回被织入增加处理目标对象
            Object target = joinPoint.getTarget();
            //为了获取注解信息
            Method currentMethod = target.getClass().getMethod(msig.getName(), msig.getParameterTypes());

            final RateLimit annotation = currentMethod.getAnnotation(RateLimit.class);
            final int limitNum = annotation.limitNum();//获取注解每秒加入桶中的token

            String functionName = msig.getName(); // 注解所在方法名区分不同的限流策略


            if (map.containsKey(functionName)) {
                rateLimiter = map.get(functionName);
            } else {
                RateLimiter rateLimiter = RateLimiter.create(limitNum);
                map.put(functionName, rateLimiter);
                this.rateLimiter = rateLimiter;
            }

            try {
                if (rateLimiter.tryAcquire()) {
                    obj = joinPoint.proceed();
                } else {
                    ResultMap result = ApiResultUtil.failResult("500", "系统繁忙");
                    final String msg = FastJsonUtils.objectToJSON(result);
                    log.info("拒绝了请求：" + msg);
                    outErrorResult(msg);
                }
            } catch (Exception e) {
                e.printStackTrace();
            } catch (Throwable throwable) {
                throwable.printStackTrace();
            }
            return obj;
        }
    }

    //将结果返回
    public void outErrorResult(String result) {
        response.setContentType("application/json;charset=UTF-8");
        try (ServletOutputStream outputStream = response.getOutputStream()) {
            outputStream.write(result.getBytes("utf-8"));
        } catch (IOException e) {
            e.printStackTrace();
        }
    }

}